from __future__ import print_function
import numpy as np
import tensorflow as tf


# 下面是一些调用的函数，后面应该放在一个utils.py中
def nhwc_to_nchw(x):
    return tf.transpose(x, [0, 3, 1, 2])
def nchw_to_nhwc(x):
    return tf.transpose(x, [0, 2, 3, 1])
def to_nhwc(image, data_format):
    new_image = nchw_to_nhwc(image)
    return new_image
def denorm_img(norm, data_format):
    return tf.clip_by_value(to_nhwc((norm + 1)*127.5, data_format), 0, 255)
def grad(img):
	kernel = tf.constant([[1 / 8, 1 / 8, 1 / 8], [1 / 8, -1, 1 / 8], [1 / 8, 1 / 8, 1 / 8]])
	kernel = tf.expand_dims(kernel, axis = -1)
	kernel = tf.expand_dims(kernel, axis = -1)
	g = tf.nn.conv2d(img, kernel, strides = [1, 1, 1, 1], padding = 'SAME')
	return g

def gradient(input):
    # filter_x=tf.reshape(tf.constant([[-1.,0.,1.],[-1.,0.,1.],[-1.,0.,1.]]),[3,3,1,1])
    # filter_y=tf.reshape(tf.constant([[-1.,-1.,-1],[0,0,0],[1,1,1]]),[3,3,1,1])
    # d_x=tf.nn.conv2d(input,filter_x,strides=[1,1,1,1], padding='SAME')
    # d_y=tf.nn.conv2d(input,filter_y,strides=[1,1,1,1], padding='SAME')
    # d=tf.sqrt(tf.square(d_x)+tf.square(d_y))
    filter = tf.reshape(tf.constant([[0., 1., 0.], [1., -4., 1.], [0., 1., 0.]]), [3, 3, 1, 1])
    d = tf.nn.conv2d(input, filter, strides=[1, 1, 1, 1], padding='SAME')
    # print(d)
    return d

def _tf_fspecial_gauss(size, sigma=1.5):
    """Function to mimic the 'fspecial' gaussian MATLAB function"""
    x_data, y_data = np.mgrid[-size//2 + 1:size//2 + 1, -size//2 + 1:size//2 + 1]

    x_data = np.expand_dims(x_data, axis=-1)
    x_data = np.expand_dims(x_data, axis=-1)

    y_data = np.expand_dims(y_data, axis=-1)
    y_data = np.expand_dims(y_data, axis=-1)

    x = tf.constant(x_data, dtype=tf.float32)
    y = tf.constant(y_data, dtype=tf.float32)

    g = tf.exp(-((x**2 + y**2)/(2.0*sigma**2)))
    return g / tf.reduce_sum(g)


def ssim_fn(img1, img2, k1=0.01, k2=0.02, L=1, window_size=11):
    """
    The function is to calculate the ssim score
    """


    window = _tf_fspecial_gauss(window_size)

    mu1 = tf.nn.conv2d(img1, window, strides = [1, 1, 1, 1], padding = 'VALID')
    mu2 = tf.nn.conv2d(img2, window, strides = [1, 1, 1, 1], padding = 'VALID')

    mu1_sq = mu1 * mu1
    mu2_sq = mu2 * mu2
    mu1_mu2 = mu1 * mu2

    sigma1_sq = tf.nn.conv2d(img1*img1, window, strides = [1 ,1, 1, 1], padding = 'VALID') - mu1_sq
    sigma2_sq = tf.nn.conv2d(img2*img2, window, strides = [1, 1, 1, 1], padding = 'VALID') - mu2_sq
    sigma1_2 = tf.nn.conv2d(img1*img2, window, strides = [1, 1, 1, 1], padding = 'VALID') - mu1_mu2

    c1 = (k1*L)**2
    c2 = (k2*L)**2

    ssim_map = ((2*mu1_mu2 + c1)*(2*sigma1_2 + c2)) / ((mu1_sq + mu2_sq + c1)*(sigma1_sq + sigma2_sq + c2))

    return tf.reduce_mean(ssim_map)










def mkdir(dir):
    import os
    if not os.path.exists(dir):
        os.makedirs(dir)






class WPrint():
    def __init__(self):
        self.file_name = "ReadableRecord.txt"
        self.cont = 0
        self.s = ""

        with open(self.file_name, "w", encoding="utf8") as f:
            pass
    def __call__(self, *s):
        s = " ".join(map(str,s))
        print(s)
        self.s = self.s + s + "\n"
        self.cont += 1
        # 25个输入写一次
        if self.cont % 15 == 0:
            with open(self.file_name, "a", encoding="utf8") as f:
                f.write(self.s)
            self.s = ""
    def finished(self):
        with open(self.file_name, "a", encoding="utf8") as f:
            f.write(self.s)



if __name__ == '__main__':
    w_print = WPrint()
    for i in range(1005):
        w_print(str(i),1)

    w_print.finished()

    # 生成dir
    mkdir("1/2/")


